{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8c3403a6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# ================================\n",
    "\n",
    "# Each user is responsible for checking the content of datasets and the\n",
    "# applicable licenses and determining if suitable for the intended use."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03166488-1651-4025-84ed-4e9e5db34933",
   "metadata": {},
   "source": [
    "<img src=\"https://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_merlin_02-deploying-multi-stage-recsys-with-merlin-systems/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "## Deploying a Multi-Stage RecSys into Production with Merlin Systems and Triton Inference Server\n",
    "\n",
    "This notebook is created using the latest stable [merlin-tensorflow](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-tensorflow/tags) container. \n",
    "\n",
    "At this point, when you reach to this notebook, we expect that you have already executed the first notebook `01-Build-RecSys-with-Merlin-Milvus.ipynb` and exported all the required files and models. \n",
    "\n",
    "We are going to generate recommended items for a given user query (user_id) by following the steps described in the figure below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38d75184-cd24-4fe3-90f4-d76028626576",
   "metadata": {},
   "source": [
    "![tritonensemble](../images/triton_ensemble.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da9dadb5-6eec-4a1b-99f9-929523f5cc07",
   "metadata": {},
   "source": [
    "Merlin Systems library have the set of operators to help serve multi-stage recommender systems built with Tensorflow on [Triton Inference Server](https://github.com/triton-inference-server/server)(TIS) easily and efficiently. Below, we will go through these operators and demonstrate their usage in serving a multi-stage system on Triton. Our workflow also includes use of the Milvus library and server for indexing vector embeddings and querying the most similar/relevant items for candidate retrieval phase."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "538677a3-acc6-48f6-acb6-d5bb5fe2e2d2",
   "metadata": {},
   "source": [
    "### Import required libraries and functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a27e18d7-b3e4-481c-b69e-23193b212c56",
   "metadata": {},
   "source": [
    "At this step, we assume you already installed `feast` library when running the first notebook `01-Build-RecSys-with-Merlin-Milvus.ipynb`. \n",
    "\n",
    "In case you need to install it for running this example, execute the following script in a cell:\n",
    "```\n",
    "%pip install feast==0.31\n",
    "```\n",
    "`Milvus` requires a milvus server to be installed and listening to requests for building and querying vector database indexes. If you have not already installed `Milvus` server and the `pymilvus` API in the first notebook, execute these lines:\n",
    "```\n",
    "%pip install milvus\n",
    "%pip install pymilvus\n",
    "```\n",
    "Note that the version of Milvus used in this notebook does not support GPU acceleration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d5c20500-1f22-43cb-8f92-fff14061d51d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# %pip install feast==0.31\n",
    "# %pip install milvus\n",
    "# %pip install pymilvus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4db1b5f1-c8fa-4e03-8744-1197873c5bee",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/cudf/utils/metadata/orc_column_statistics_pb2.py:19: DeprecationWarning: Call to deprecated create function FileDescriptor(). Note: Create unlinked descriptors is going to go away. Please use get/find descriptors from generated code or query the descriptor_pool.\n",
      "  DESCRIPTOR = _descriptor.FileDescriptor(\n",
      "/usr/local/lib/python3.8/dist-packages/cudf/utils/metadata/orc_column_statistics_pb2.py:37: DeprecationWarning: Call to deprecated create function FieldDescriptor(). Note: Create unlinked descriptors is going to go away. Please use get/find descriptors from generated code or query the descriptor_pool.\n",
      "  _descriptor.FieldDescriptor(\n",
      "/usr/local/lib/python3.8/dist-packages/cudf/utils/metadata/orc_column_statistics_pb2.py:30: DeprecationWarning: Call to deprecated create function Descriptor(). Note: Create unlinked descriptors is going to go away. Please use get/find descriptors from generated code or query the descriptor_pool.\n",
      "  _INTEGERSTATISTICS = _descriptor.Descriptor(\n",
      "2023-07-02 03:09:04.730794: I tensorflow/core/platform/cpu_feature_guard.cc:183] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "/usr/local/lib/python3.8/dist-packages/merlin/dtypes/mappings/torch.py:43: UserWarning: PyTorch dtype mappings did not load successfully due to an error: No module named 'torch'\n",
      "  warn(f\"PyTorch dtype mappings did not load successfully due to an error: {exc.msg}\")\n",
      "/usr/local/lib/python3.8/dist-packages/nvtabular/loader/__init__.py:19: DeprecationWarning: The `nvtabular.loader` module has moved to a new repository, at https://github.com/NVIDIA-Merlin/dataloader .  Support for importing from `nvtabular.loader` is deprecated, and will be removed in a future version. Please update your imports to refer to `merlinloader`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"TF_GPU_ALLOCATOR\"]=\"cuda_malloc_async\"\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import feast\n",
    "import seedir as sd\n",
    "from nvtabular import ColumnSchema, Schema\n",
    "\n",
    "from merlin.systems.dag.ensemble import Ensemble\n",
    "from merlin.systems.dag.ops.session_filter import FilterCandidates\n",
    "from merlin.systems.dag.ops.softmax_sampling import SoftmaxSampling\n",
    "from merlin.systems.dag.ops.tensorflow import PredictTensorflow\n",
    "from merlin.systems.dag.ops.unroll_features import UnrollFeatures\n",
    "from merlin.systems.triton.utils import send_triton_request"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53e705dd-721c-4a0d-9d34-42b31b20e46e",
   "metadata": {},
   "source": [
    "Execute the following cell to start the milvus server (for performing embedding vector ANN search) on the local server. Set the default server base directory to any local directory of your choice (this is where milvus creates logs, indices, etc. to make the ANN search possible)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "678bbb47-7caf-494f-b63d-7889f609c0b0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "    __  _________ _   ____  ______\n",
      "   /  |/  /  _/ /| | / / / / / __/\n",
      "  / /|_/ // // /_| |/ / /_/ /\\ \\\n",
      " /_/  /_/___/____/___/\\____/___/ {Lite}\n",
      "\n",
      " Welcome to use Milvus!\n",
      "\n",
      " Version:   v2.2.10-lite\n",
      " Process:   358956\n",
      " Started:   2023-07-02 03:09:23\n",
      " Config:    /workspace/data/fstore_milvus/server/configs/milvus.yaml\n",
      " Logs:      /workspace/data/fstore_milvus/server/logs\n",
      "\n",
      " Ctrl+C to exit ...\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "from milvus import default_server\n",
    "from pymilvus import (connections, FieldSchema, CollectionSchema, DataType, Collection, utility)\n",
    "\n",
    "BASE_DIR = os.environ.get(\"BASE_DIR\", \"/workspace/data/fstore_milvus/\")\n",
    "\n",
    "default_server.set_base_dir(os.path.join(BASE_DIR, 'server'))\n",
    "default_server.cleanup()    # this will remove all collections and indexes previously created\n",
    "default_server.start()\n",
    "_HOST = '127.0.0.1'\n",
    "_PORT = default_server.listen_port\n",
    "\n",
    "fmt = \"\\n=== {:30} ===\\n\"\n",
    "search_latency_fmt = \"search latency = {:.4f}s\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61d09468-0395-457a-a2b7-6369ebaa6432",
   "metadata": {},
   "source": [
    "You can use the test script below to create a connection to the `milvus` server and test it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "45ed5c93-ca38-457a-a650-3de59e96de42",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#connections.connect(host=_HOST, port=_PORT)\n",
    "#print(connections.list_connections())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55ead20e-c573-462e-9aa2-c3494bf0129f",
   "metadata": {},
   "source": [
    "### Register our features on feature store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2ac115e-4794-4a69-a962-8481f6e86df3",
   "metadata": {},
   "source": [
    "The Feast feature registry is a central catalog of all the feature definitions and their related metadata(read more [here](https://docs.feast.dev/getting-started/architecture-and-components/registry)). We have defined our user and item features definitions in the `user_features.py` and  `item_features.py` files. With FeatureView() users can register data sources in their organizations into Feast, and then use those data sources for both training and online inference. In the `user_features.py` and `item_features.py` files, we are telling Feast where to find user and item features.\n",
    "\n",
    "Before we move on to the next steps, we need to perform `feast apply`command as directed below.  With that, we register our features, we can apply the changes to create our feature registry and store all entity and feature view definitions in a local SQLite online store called `online_store.db`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "66c02d67-df45-4869-8262-647cba77efcb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# define feature repo path\n",
    "feast_repo_path = os.path.join(BASE_DIR, \"feast_repo/feature_repo/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e5fa545b-a979-4216-b176-ffd70d66e69d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/workspace/data/fstore_milvus/feast_repo/feature_repo\n",
      "Created entity \u001b[1m\u001b[32muser_id_raw\u001b[0m\n",
      "Created entity \u001b[1m\u001b[32mitem_id\u001b[0m\n",
      "Created feature view \u001b[1m\u001b[32mitem_features\u001b[0m\n",
      "Created feature view \u001b[1m\u001b[32muser_features\u001b[0m\n",
      "\n",
      "Created sqlite table \u001b[1m\u001b[32mfeast_repo_item_features\u001b[0m\n",
      "Created sqlite table \u001b[1m\u001b[32mfeast_repo_user_features\u001b[0m\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%cd $feast_repo_path\n",
    "!feast apply"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c641fcd2-bd11-4569-80d4-2ae5e01a5cad",
   "metadata": {},
   "source": [
    "### Loading features from offline store into an online store \n",
    "\n",
    "After we executed `apply` and registered our features and created our online local store, now we need to perform [materialization](https://docs.feast.dev/how-to-guides/running-feast-in-production) operation. This is done to keep our online store up to date and get it ready for prediction. For that we need to run a job that loads feature data from our feature view sources into our online store. As we add new features to our offline stores, we can continuously materialize them to keep our online store up to date by finding the latest feature values for each user. \n",
    "\n",
    "When you run the `feast materialize ..` command below, you will see a message <i>Materializing 2 feature views from 1995-01-01 01:01:01+00:00 to 2025-01-01 01:01:01+00:00 into the sqlite online store </i>  will be printed out.\n",
    "\n",
    "Note that the materialization step takes some time... "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "52dacbbc-bdb6-4f7a-b202-3802050f0362",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Materializing \u001b[1m\u001b[32m2\u001b[0m feature views from \u001b[1m\u001b[32m1995-01-01 01:01:01+00:00\u001b[0m to \u001b[1m\u001b[32m2025-01-01 01:01:01+00:00\u001b[0m into the \u001b[1m\u001b[32msqlite\u001b[0m online store.\n",
      "\n",
      "\u001b[1m\u001b[32mitem_features\u001b[0m:\n",
      "100%|███████████████████████████████████████████████████████| 49008/49008 [00:17<00:00, 2742.34it/s]\n",
      "\u001b[1m\u001b[32muser_features\u001b[0m:\n",
      "100%|█████████████████████████████████████████████████| 7305761/7305761 [1:08:59<00:00, 1764.94it/s]\n"
     ]
    }
   ],
   "source": [
    "!feast materialize 1995-01-01T01:01:01 2025-01-01T01:01:01"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fcc26e6-f6f3-4e44-bf3c-3b8e66dc9fd6",
   "metadata": {},
   "source": [
    "Now, let's check our `feature_repo` structure again after we ran `apply` and `materialize` commands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9caba4e3-e6e0-4e2f-b51d-cd3456fd4a63",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "feast_repo/\n",
      "├─README.md\n",
      "├─__init__.py\n",
      "└─feature_repo/\n",
      "  ├─__init__.py\n",
      "  ├─data/\n",
      "  │ ├─item_features.parquet\n",
      "  │ ├─online_store.db\n",
      "  │ ├─registry.db\n",
      "  │ └─user_features.parquet\n",
      "  ├─feature_store.yaml\n",
      "  ├─item_features.py\n",
      "  ├─test_workflow.py\n",
      "  └─user_features.py\n"
     ]
    }
   ],
   "source": [
    "sd.seedir(\n",
    "    os.path.join(BASE_DIR, \"feast_repo\"),\n",
    "    style=\"lines\",\n",
    "    itemlimit=10,\n",
    "    depthlimit=5,\n",
    "    exclude_folders=['.ipynb_checkpoints', '__pycache__'],\n",
    "    sort=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e768637c-0a4d-404b-8b58-7182fef0ab0e",
   "metadata": {},
   "source": [
    "### Set up Milvus index, create feature store client and objects for the Triton ensemble"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa037c0-7dad-427c-98bb-3da413e8fd14",
   "metadata": {},
   "source": [
    "To set up an index with `milvus`, you need to retrieve item embeddings and build a `milvus` collection with them. The function `setup_milvus` below will take as input an item embeddings vector `item_vector` along with the `milvus` server connection parameters (`HOST` and `PORT`) and a name `coll_name` for the collection to build. The function `query_milvus` will take as input a user embedding vector and query it against the index built by `setup_milvus`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c8ce5aa9-77ef-4b46-a6f5-56e15ac492b3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def setup_milvus(item_vector, HOST, PORT, coll_name):\n",
    "    \"\"\"\n",
    "    Function that will create a Milvus index from an embedding vector. Currently only supports L2 distance.\n",
    "    Parameters\n",
    "    ----------\n",
    "    item_vector : Numpy.ndarray\n",
    "        This is a matrix representing all the nodes embeddings, represented as a numpy ndarray.\n",
    "    HOST : str\n",
    "        milvus server host ip address\n",
    "    PORT : integer\n",
    "        milvus server listening port number\n",
    "    coll_name : string\n",
    "        name of milvus collection to create and store/index item embeddings\n",
    "    \"\"\"\n",
    "    \n",
    "    n_dim = item_vector.shape[1]-1     # size of the embedding vector - 1 (first value is item_id)\n",
    "\n",
    "    #######################################################################################\n",
    "    # 1. connect to Milvus\n",
    "    # Add a new connection alias `default` for Milvus server in `localhost:19530`\n",
    "    # Actually the \"default\" alias is a buildin in PyMilvus.\n",
    "    # If the address of Milvus is the same as `localhost:19530`, you can omit all\n",
    "    # parameters and call the method as: `connections.connect()`.\n",
    "    # Note: the `using` parameter of the following methods is default to \"default\".\n",
    "\n",
    "    print(fmt.format(\"Connecting to Milvus...\"))\n",
    "    connections.connect(\"default\", host=HOST, port=PORT)\n",
    "    has = utility.has_collection(coll_name)\n",
    "    print(f\"Does collection {coll_name} exist in Milvus: {has}\")\n",
    "    \n",
    "    #######################################################################################\n",
    "    # 2. create collection\n",
    "    # We're going to create a collection with 2 fields.\n",
    "    # +-+-------------------+-------------+------------------+----------------------------+\n",
    "    # | |    field name     | field type  | other attributes |      field description     |\n",
    "    # +-+-------------------+-------------+------------------+----------------------------+\n",
    "    # |1|       \"pk\"        |    Int64    |  is_primary=True |      \"primary key\"         |\n",
    "    # | |                   |             |   auto_id=False  |                            |\n",
    "    # +-+-------------------+-------------+------------------+----------------------------+\n",
    "    # |2| \"item_embeddings\" | FloatVector |     dim=64       | \"float vector with dim 64\" |\n",
    "    # +-+-------------------+-------------+------------------+----------------------------+\n",
    "\n",
    "    fields = [\n",
    "        FieldSchema(name=\"pk\", dtype=DataType.INT64, is_primary=True, auto_id=False),\n",
    "        FieldSchema(name=\"item_embeddings\", dtype=DataType.FLOAT_VECTOR, dim=n_dim)\n",
    "    ]\n",
    "    schema = CollectionSchema(fields, \"fields that represent RecSys item embeddings.\")\n",
    "\n",
    "    print(fmt.format(\"Creating collection...\"))\n",
    "    coll_milvus = Collection(coll_name, schema, consistency_level=\"Strong\")\n",
    "    print(f\"Collection created.\")\n",
    "    \n",
    "    ######################################################################################\n",
    "    # 3. insert data\n",
    "    # We are going to insert item embeddings into `coll_milvus`\n",
    "    # Data to be inserted must be organized in fields.\n",
    "    #\n",
    "    # The insert() method returns:\n",
    "    # - either automatically generated primary keys by Milvus if auto_id=True in the schema;\n",
    "    # - or the existing primary key field from the entities if auto_id=False in the schema.\n",
    "\n",
    "    print(fmt.format(\"Inserting embeddings...\"))\n",
    "    num_entities = item_vector.shape[0]\n",
    "    entities = [\n",
    "        # provide the pk field, which is item id, because `auto_id` is set to False\n",
    "        [item_vector[i][0].astype(np.int64) for i in range(num_entities)],\n",
    "        [[item_vector[i][j+1] for j in range(n_dim)] for i in range(num_entities)],  # field embeddings\n",
    "    ]\n",
    "    insert_result = coll_milvus.insert(entities)\n",
    "    coll_milvus.flush()\n",
    "    print(f\"Number of entities in Milvus: {coll_milvus.num_entities}\")  # check number inserted\n",
    "\n",
    "    #######################################################################################\n",
    "    # 4. create index\n",
    "    # We are going to create an IVF_FLAT index for coll_milvus collection.\n",
    "    # create_index() can only be applied to `FloatVector` and `BinaryVector` fields.\n",
    "    print(fmt.format(\"Creating index IVF_FLAT...\"))\n",
    "    index = {\n",
    "        \"index_type\": \"IVF_FLAT\",\n",
    "        \"metric_type\": \"L2\",\n",
    "        \"params\": {\"nlist\": 128},\n",
    "    }\n",
    "    start_time = time.time()\n",
    "    coll_milvus.create_index(\"item_embeddings\", index)   # will save the index to a default location of milvus installation\n",
    "    end_time = time.time()\n",
    "    print(f\"Indexing latency: {end_time - start_time}\")\n",
    "\n",
    "def query_milvus(user_vector, coll_name, topk=10):\n",
    "    \"\"\"\n",
    "    Function to query input user vector against an indexed set item embedding vectors.\n",
    "    Parameters\n",
    "    ----------\n",
    "    user_vector : Numpy.ndarray\n",
    "        This is a 2D matrix representing the user embedding vectors to be queried, represented as a numpy ndarray\n",
    "    coll_name : str\n",
    "        Name of milvus collection where item embeddings are stored and indexed\n",
    "    topk : integer\n",
    "        number of most similar items to return\n",
    "    \"\"\"\n",
    "    print(fmt.format(\"Loading collection and index...\"))\n",
    "    coll_milvus = Collection(coll_name)      # Get an existing collection.    \n",
    "    coll_milvus.load()\n",
    "    \n",
    "    print(fmt.format(\"Searching based on vector similarity...\"))\n",
    "    search_params = {\n",
    "        \"metric_type\": \"L2\",\n",
    "        \"params\": {\"nprobe\": 10},\n",
    "    }\n",
    "    \n",
    "    start_time = time.time()\n",
    "    result = coll_milvus.search(user_vector, \"item_embeddings\", search_params, limit=topk, output_fields=[\"pk\"])\n",
    "    end_time = time.time()\n",
    "\n",
    "    i=0\n",
    "    for hits in result:\n",
    "        i=i+1; print(f\"Query for user {i}:\")\n",
    "        for hit in hits:\n",
    "            print(f\"hit: {hit}, pk: {hit.entity.get('pk')}\")\n",
    "\n",
    "    print(search_latency_fmt.format(end_time - start_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2c172b4-16a7-4c7b-b823-6f30ca467713",
   "metadata": {},
   "source": [
    "Now let's use these two functions to setup the milvus item embeddings vector index and query similarity between one user's embedding against the item embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a39bac5f-0078-4a05-ba6f-d3fbd3fa2236",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Connecting to Milvus...        ===\n",
      "\n",
      "Does collection coll_item exist in Milvus: False\n",
      "\n",
      "=== Creating collection...         ===\n",
      "\n",
      "Collection created.\n",
      "\n",
      "=== Inserting embeddings...        ===\n",
      "\n",
      "Number of entities in Milvus: 49008\n",
      "\n",
      "=== Creating index IVF_FLAT...     ===\n",
      "\n",
      "Indexing latency: 3.735898733139038\n"
     ]
    }
   ],
   "source": [
    "item_embeddings = np.ascontiguousarray(\n",
    "    pd.read_parquet(os.path.join(BASE_DIR, \"item_embeddings.parquet\")).to_numpy()\n",
    ")\n",
    "setup_milvus(item_embeddings, _HOST, _PORT, \"coll_item\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "731fb049-b016-40b8-8cf6-5dbbb92ea4fd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Loading collection and index... ===\n",
      "\n",
      "\n",
      "=== Searching based on vector similarity... ===\n",
      "\n",
      "Query for user 1:\n",
      "hit: id: 27545, distance: 104.72942352294922, entity: {'pk': 27545}, pk: 27545\n",
      "hit: id: 30332, distance: 105.61385345458984, entity: {'pk': 30332}, pk: 30332\n",
      "hit: id: 36994, distance: 106.31818389892578, entity: {'pk': 36994}, pk: 36994\n",
      "hit: id: 40632, distance: 106.88359069824219, entity: {'pk': 40632}, pk: 40632\n",
      "hit: id: 37987, distance: 106.98526000976562, entity: {'pk': 37987}, pk: 37987\n",
      "hit: id: 24961, distance: 107.1307373046875, entity: {'pk': 24961}, pk: 24961\n",
      "hit: id: 42383, distance: 107.14056396484375, entity: {'pk': 42383}, pk: 42383\n",
      "hit: id: 42941, distance: 107.36529541015625, entity: {'pk': 42941}, pk: 42941\n",
      "hit: id: 36093, distance: 107.47981262207031, entity: {'pk': 36093}, pk: 36093\n",
      "hit: id: 35356, distance: 107.65443420410156, entity: {'pk': 35356}, pk: 35356\n",
      "Query for user 2:\n",
      "hit: id: 36206, distance: 34.583248138427734, entity: {'pk': 36206}, pk: 36206\n",
      "hit: id: 47395, distance: 35.64282989501953, entity: {'pk': 47395}, pk: 47395\n",
      "hit: id: 34507, distance: 35.77540588378906, entity: {'pk': 34507}, pk: 34507\n",
      "hit: id: 42040, distance: 35.929046630859375, entity: {'pk': 42040}, pk: 42040\n",
      "hit: id: 43496, distance: 36.14706039428711, entity: {'pk': 43496}, pk: 43496\n",
      "hit: id: 35356, distance: 36.21003723144531, entity: {'pk': 35356}, pk: 35356\n",
      "hit: id: 36994, distance: 36.278682708740234, entity: {'pk': 36994}, pk: 36994\n",
      "hit: id: 38036, distance: 36.30451202392578, entity: {'pk': 38036}, pk: 38036\n",
      "hit: id: 29023, distance: 36.41819763183594, entity: {'pk': 29023}, pk: 29023\n",
      "hit: id: 43565, distance: 36.41949462890625, entity: {'pk': 43565}, pk: 43565\n",
      "search latency = 0.3253s\n"
     ]
    }
   ],
   "source": [
    "user_embeddings = np.ascontiguousarray(\n",
    "    pd.read_parquet(os.path.join(BASE_DIR, \"user_embeddings.parquet\")).to_numpy()\n",
    ")\n",
    "query_milvus(user_embeddings[0:2,1:], \"coll_item\", 10)  # query embedding vectors for the first 2 users, ignoring first vector value which is user_id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb0ea6ab-eb38-4799-b6d0-61fcff0e1660",
   "metadata": {},
   "source": [
    "Now define paths for ranking model and retrieval model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "23ba59b5-08c3-44b5-86f2-e63dec6893af",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "retrieval_model_path = os.path.join(BASE_DIR, \"query_tower/\")\n",
    "ranking_model_path = os.path.join(BASE_DIR, \"dlrm/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b996019-bd2a-44e0-b004-4f412b300d63",
   "metadata": {},
   "source": [
    "The custome-written `QueryMilvus` operator defined below creates an interface between a `Milvus` IVF_FLAT index and Triton Inference Server. For a given input query vector, we do a nearest neighbor search query to find the ids of top-k nearby embeddings in the index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a30ef757-60ef-4f38-989a-410af8711720",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import pathlib\n",
    "from pathlib import Path\n",
    "from typing import Dict, List\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from merlin.core.dispatch import HAS_GPU\n",
    "from merlin.core.protocols import Transformable\n",
    "from merlin.dag import ColumnSelector\n",
    "from merlin.schema import ColumnSchema, Schema\n",
    "from merlin.systems.dag.ops.operator import InferenceOperator\n",
    "\n",
    "import pymilvus\n",
    "\n",
    "class QueryMilvus(InferenceOperator):\n",
    "    \"\"\"\n",
    "    This operator creates an interface between a Milvus Approximate Nearest Neighbors (ANN)\n",
    "    Index and Triton Infrence Server. The operator allows users to perform different supported\n",
    "    types of Nearest Neighbor search to your ensemble. For input query vector, we do an ANN\n",
    "    search query to find the ids of top-k nearby nodes in the index.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, milvus_host: str, milvus_port: str, collection_name: str, index_name: str, topk=10):\n",
    "        \"\"\"\n",
    "        Creates a QueryMilvus Pipelineable Inference Operator.\n",
    "        Parameters\n",
    "        ----------\n",
    "        milvus_host: str\n",
    "            Milvus host name (eg. 127.0.0.1)\n",
    "        milvus_port: str\n",
    "            Milvus port number (eg. 19530)\n",
    "        collection_name: str\n",
    "            Milvus collection name where item embeddings have been previously stored and indexed\n",
    "        index_name: str\n",
    "            Name of the milvus index for item embeddings\n",
    "        topk : int, optional\n",
    "            The number of results we should receive from query to Milvus as output, by default 10\n",
    "        \"\"\"\n",
    "        print(\"QueryMilvus operator initializing...\")\n",
    "        super().__init__()\n",
    "        self.mlv_host = milvus_host\n",
    "        self.mlv_port = milvus_port\n",
    "        self.col_name = collection_name\n",
    "        self.idx_name = index_name\n",
    "        self.topk = topk\n",
    "\n",
    "    def load_artifacts(self, artifact_path):\n",
    "        \"\"\"\n",
    "        Makes a connection to the milvus server and loads the indexed embeddings collection\n",
    "        Parameters\n",
    "        ----------\n",
    "        artifact_path: str\n",
    "            Unused - pass empty string\n",
    "        \"\"\"\n",
    "        self.mlv_con = pymilvus.connections\n",
    "        self.mlv_con.connect(host=self.mlv_host, port=self.mlv_port)\n",
    "        self.mlv_coll = Collection(self.col_name)      # Get the right milvus collection.    \n",
    "        self.mlv_coll.load()\n",
    "\n",
    "    def transform(\n",
    "        self, col_selector: ColumnSelector, transformable: Transformable\n",
    "    ) -> Transformable:\n",
    "        \"\"\"\n",
    "        Transform input dataframe to output dataframe using function logic.\n",
    "        In the case of QueryMilvus operator, it takes the input user embeddings vector,\n",
    "        queries the most similar k item embeddings and returns their id's.\n",
    "        \"\"\"\n",
    "        print(fmt.format(\"QueryMilvus operator transforming...\"))\n",
    "        user_vector = list(transformable.values())[0]\n",
    "        user_vector = user_vector.values\n",
    "\n",
    "        print(fmt.format(\"Searching based on vector similarity...\"))\n",
    "        search_params = {\n",
    "            \"metric_type\": \"L2\",\n",
    "            \"params\": {\"nprobe\": 10},\n",
    "        }\n",
    "\n",
    "        start_time = time.time()\n",
    "        result = self.mlv_coll.search(user_vector, self.idx_name, search_params, limit=self.topk, output_fields=[\"pk\"])\n",
    "        end_time = time.time()\n",
    "        print(f\"Search latency: {end_time - start_time}\")\n",
    "        \n",
    "        candidate_ids = []\n",
    "        for hits in result:\n",
    "            for hit in hits:\n",
    "                candidate_ids.append(hit.entity.get('pk'))\n",
    "        candidate_ids = np.array(candidate_ids).T.astype(np.int32)\n",
    "\n",
    "        return type(transformable)({\"candidate_ids\": candidate_ids})\n",
    "\n",
    "    def compute_input_schema(\n",
    "        self,\n",
    "        root_schema: Schema,\n",
    "        parents_schema: Schema,\n",
    "        deps_schema: Schema,\n",
    "        selector: ColumnSelector,\n",
    "    ) -> Schema:\n",
    "        \"\"\"\n",
    "        Compute the input schema of this node given the root, parents and dependencies schemas of\n",
    "        all ancestor nodes.\n",
    "        Parameters\n",
    "        ----------\n",
    "        root_schema : Schema\n",
    "            The schema representing the input columns to the graph\n",
    "        parents_schema : Schema\n",
    "            A schema representing all the output columns of the ancestors of this node.\n",
    "        deps_schema : Schema\n",
    "            A schema representing the dependencies of this node.\n",
    "        selector : ColumnSelector\n",
    "            A column selector representing a target subset of columns necessary for this node's\n",
    "            operator\n",
    "        Returns\n",
    "        -------\n",
    "        Schema\n",
    "            A schema that has the correct representation of all the incoming columns necessary for\n",
    "            this node's operator to complete its transform.\n",
    "        Raises\n",
    "        ------\n",
    "        ValueError\n",
    "            Cannot receive more than one input for this node\n",
    "        \"\"\"\n",
    "        input_schema = super().compute_input_schema(\n",
    "            root_schema, parents_schema, deps_schema, selector\n",
    "        )\n",
    "        return input_schema\n",
    "\n",
    "    def compute_output_schema(\n",
    "        self, input_schema: Schema, col_selector: ColumnSelector, prev_output_schema: Schema = None\n",
    "    ) -> Schema:\n",
    "        \"\"\"\n",
    "        Compute the input schema of this node given the root, parents and dependencies schemas of\n",
    "        all ancestor nodes.\n",
    "        Parameters\n",
    "        ----------\n",
    "        input_schema : Schema\n",
    "            The schema representing the input columns to the graph\n",
    "        col_selector : ColumnSelector\n",
    "            A column selector representing a target subset of columns necessary for this node's\n",
    "            operator\n",
    "        prev_output_schema : Schema\n",
    "            A schema representing the output of the previous node.\n",
    "        Returns\n",
    "        -------\n",
    "        Schema\n",
    "            A schema object representing all outputs of this node.\n",
    "        \"\"\"\n",
    "        return Schema(\n",
    "            [\n",
    "                ColumnSchema(\"candidate_ids\", dtype=np.int32),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def validate_schemas(\n",
    "        self, parents_schema, deps_schema, input_schema, output_schema, strict_dtypes=False\n",
    "    ):\n",
    "        if len(input_schema.column_schemas) > 1:\n",
    "            raise ValueError(\n",
    "                \"More than one input has been detected for this node,\"\n",
    "                / f\"inputs received: {input_schema.column_names}\"\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46697177-512a-473e-8cca-9fe51d3daa03",
   "metadata": {},
   "source": [
    "Now create the feature store client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3bc00e04-c70c-4882-9952-66f4dbb97bdc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "feature_store = feast.FeatureStore(feast_repo_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c45df06-0cbe-4b52-ac1f-786e763895d7",
   "metadata": {},
   "source": [
    "Fetch user features with `QueryFeast` operator from the feature store. `QueryFeast` operator is responsible for ensuring that our feast feature store can communicate correctly with tritonserver for the ensemble feast feature look ups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3decbe7b-03e3-4978-baac-03f6a0b078c9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Materializing \u001b[1m\u001b[32m1\u001b[0m feature views to \u001b[1m\u001b[32m2023-07-02 12:25:26+00:00\u001b[0m into the \u001b[1m\u001b[32msqlite\u001b[0m online store.\n",
      "\n",
      "\u001b[1m\u001b[32muser_features\u001b[0m from \u001b[1m\u001b[32m2025-01-01 01:01:01+00:00\u001b[0m to \u001b[1m\u001b[32m2023-07-02 12:25:26+00:00\u001b[0m:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]\n"
     ]
    }
   ],
   "source": [
    "from merlin.systems.dag.ops.feast import QueryFeast \n",
    "\n",
    "user_features = [\"user_id_raw\"] >> QueryFeast.from_feature_view(\n",
    "    store=feature_store,\n",
    "    view=\"user_features\",\n",
    "    column=\"user_id_raw\",\n",
    "    include_id=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27e25be7-3ff0-49c2-a3fc-03ec4d615e77",
   "metadata": {},
   "source": [
    "Retrieve top-K candidate items using `retrieval model` that are relevant for a given user. We use `PredictTensorflow()` operator that takes a tensorflow model and packages it correctly for TIS to run with the tensorflow backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "21139caa-3a51-42e6-b006-21a92c95f1bc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function tensorflow.python.dlpack.dlpack.from_dlpack(dlcapsule)>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# prevent TF to claim all GPU memory\n",
    "from merlin.dataloader.tf_utils import configure_tensorflow\n",
    "\n",
    "configure_tensorflow()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "47c2d9b1-51dc-4549-977d-d7941ee6486c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-07-02 12:29:33.157366: I tensorflow/core/common_runtime/gpu/gpu_process_state.cc:226] Using CUDA malloc Async allocator for GPU: 0\n",
      "2023-07-02 12:29:33.157617: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1638] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 16255 MB memory:  -> device: 0, name: Tesla V100-SXM2-32GB-LS, pci bus id: 0000:85:00.0, compute capability: 7.0\n",
      "WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 28). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: /tmp/tmp98qi4gl1/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: /tmp/tmp98qi4gl1/assets\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QueryMilvus operator initializing...\n"
     ]
    }
   ],
   "source": [
    "topk_retrieval = int(\n",
    "    os.environ.get(\"topk_retrieval\", \"100\")\n",
    ")\n",
    "retrieval = (\n",
    "    user_features\n",
    "    >> PredictTensorflow(retrieval_model_path)\n",
    "    >> QueryMilvus(\"127.0.0.1\", str(default_server.listen_port), \"coll_item\", \"item_embeddings\", topk=topk_retrieval)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ce4429c-1fe1-4304-bcdf-badebe3b5485",
   "metadata": {},
   "source": [
    "Fetch item features for the candidate items that are retrieved from the retrieval step above from the feature store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b270f663-0ae1-4356-acd4-5f8c986abf4d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): usage.feast.dev:443\n",
      "DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): usage.feast.dev:443\n",
      "DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): usage.feast.dev:443\n",
      "DEBUG:urllib3.connectionpool:https://usage.feast.dev:443 \"POST / HTTP/1.1\" 200 2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Materializing \u001b[1m\u001b[32m1\u001b[0m feature views to \u001b[1m\u001b[32m2023-07-02 12:30:07+00:00\u001b[0m into the \u001b[1m\u001b[32msqlite\u001b[0m online store.\n",
      "\n",
      "\u001b[1m\u001b[32mitem_features\u001b[0m from \u001b[1m\u001b[32m2025-01-01 01:01:01+00:00\u001b[0m to \u001b[1m\u001b[32m2023-07-02 12:30:07+00:00\u001b[0m:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): usage.feast.dev:443\n",
      "DEBUG:urllib3.connectionpool:https://usage.feast.dev:443 \"POST / HTTP/1.1\" 200 2\n",
      "DEBUG:urllib3.connectionpool:https://usage.feast.dev:443 \"POST / HTTP/1.1\" 200 2\n",
      "DEBUG:urllib3.connectionpool:https://usage.feast.dev:443 \"POST / HTTP/1.1\" 200 2\n",
      "0it [00:00, ?it/s]\n",
      "DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): usage.feast.dev:443\n",
      "DEBUG:urllib3.connectionpool:https://usage.feast.dev:443 \"POST / HTTP/1.1\" 200 2\n"
     ]
    }
   ],
   "source": [
    "item_features = retrieval[\"candidate_ids\"] >> QueryFeast.from_feature_view(\n",
    "    store=feature_store,\n",
    "    view=\"item_features\",\n",
    "    column=\"candidate_ids\",\n",
    "    output_prefix=\"item\",\n",
    "    include_id=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "304a4d09-db05-4666-b520-75dbbbc7ab17",
   "metadata": {},
   "source": [
    "Merge the user features and items features to create the all set of combined features that were used in model training using `UnrollFeatures` operator which takes a target column and joins the \"unroll\" columns to the target. This helps when broadcasting a series of user features to a set of items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "eb0ef434-03a5-4a36-afb9-e19a43243c64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "user_features_to_unroll = [\n",
    "    \"user_id\",\n",
    "    \"user_age\",\n",
    "]\n",
    "\n",
    "combined_features = item_features >> UnrollFeatures(\n",
    "    \"item_id\", user_features[user_features_to_unroll]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fb0ce66-6b6c-43be-885e-a5435c3bbd9e",
   "metadata": {},
   "source": [
    "Rank the combined features using the trained ranking model, which is a DLRM model for this example. We feed the path of the ranking model to `PredictTensorflow()` operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ce31723e-af4d-4827-bb60-3a9fafcd9da6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 50). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: /tmp/tmpfrg2zgoo/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: /tmp/tmpfrg2zgoo/assets\n"
     ]
    }
   ],
   "source": [
    "ranking = combined_features >> PredictTensorflow(ranking_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f86fa47-de61-4007-ab55-9076e12ce963",
   "metadata": {},
   "source": [
    "For the ordering we use `SoftmaxSampling()` operator. This operator sorts all inputs in descending order given the input ids and prediction introducing some randomization into the ordering by sampling items from the softmax of the predicted relevance scores, and finally returns top-k ordered items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "7f65598b-e3e7-4238-a73e-19d00c3deb26",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "top_k=10\n",
    "ordering = combined_features[\"item_id_raw\"] >> SoftmaxSampling(\n",
    "    relevance_col=ranking[\"click/binary_classification_task\"], topk=top_k, temperature=20.0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4e2e389-d884-44a1-8e32-4916a0eb43cf",
   "metadata": {},
   "source": [
    "### Export Graph as Ensemble\n",
    "The last step is to create the ensemble artifacts that TIS can consume. To make these artifacts, we import the Ensemble class. This class  represents an entire ensemble consisting of multiple models that run sequentially in TIS initiated by an inference request. It is responsible for interpreting the graph and exporting the correct files for TIS.\n",
    "\n",
    "When we create an Ensemble object we feed the graph and a schema representing the starting input of the graph.  After we create the ensemble object, we export the graph, supplying an export path for the `ensemble.export()` function. This returns an ensemble config which represents the entire inference pipeline and a list of node-specific configs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50bc2e4f-5e58-4ad4-8ae5-d79ad286978f",
   "metadata": {},
   "source": [
    "Create the folder to export the models and config files (remove the folder first, if it exists)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b28c452f-543c-45a4-9995-130ca6919669",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if os.path.isdir(os.path.join(BASE_DIR, 'poc_ensemble')):\n",
    "    !rm -r {os.path.join(BASE_DIR, 'poc_ensemble')}\n",
    "os.makedirs(os.path.join(BASE_DIR, 'poc_ensemble'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a061bd82-e553-4823-8d14-3ae88a458c14",
   "metadata": {},
   "source": [
    "Create a request schema that we are going to use when sending a request to Triton Inference Server (TIS)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "9c8b7b94-5559-4587-a272-4d9de2d53dd1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "request_schema = Schema(\n",
    "    [\n",
    "        ColumnSchema(\"user_id_raw\", dtype=np.int32),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6c64d686-aed5-42f8-b517-482b4237c69f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['ordered_ids', 'ordered_scores']\n"
     ]
    }
   ],
   "source": [
    "# define the path where all the models and config files exported to\n",
    "export_path = os.path.join(BASE_DIR, 'poc_ensemble')\n",
    "\n",
    "ensemble = Ensemble(ordering, request_schema)\n",
    "ens_config, node_configs = ensemble.export(export_path)\n",
    "\n",
    "# return the output column name\n",
    "outputs = ensemble.graph.output_schema.column_names\n",
    "print(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4b73fee-abaa-4712-8f4b-d08bb7ae267b",
   "metadata": {},
   "source": [
    "Let's check our `export_path` structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "22665010-6d4e-45da-afcd-4b98b0f0ebaa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "poc_ensemble/\n",
      "├─0_predicttensorflowtriton/\n",
      "│ ├─1/\n",
      "│ │ └─model.savedmodel/\n",
      "│ │   ├─assets/\n",
      "│ │   ├─fingerprint.pb\n",
      "│ │   ├─keras_metadata.pb\n",
      "│ │   ├─saved_model.pb\n",
      "│ │   └─variables/\n",
      "│ │     ├─variables.data-00000-of-00001\n",
      "│ │     └─variables.index\n",
      "│ └─config.pbtxt\n",
      "├─1_predicttensorflowtriton/\n",
      "│ ├─1/\n",
      "│ │ └─model.savedmodel/\n",
      "│ │   ├─.merlin/\n",
      "│ │   │ ├─input_schema.json\n",
      "│ │   │ └─output_schema.json\n",
      "│ │   ├─assets/\n",
      "│ │   ├─fingerprint.pb\n",
      "│ │   ├─keras_metadata.pb\n",
      "│ │   ├─saved_model.pb\n",
      "│ │   └─variables/\n",
      "│ │     ├─variables.data-00000-of-00001\n",
      "│ │     └─variables.index\n",
      "│ └─config.pbtxt\n",
      "└─executor_model/\n",
      "  ├─1/\n",
      "  │ ├─ensemble/\n",
      "  │ │ ├─ensemble.pkl\n",
      "  │ │ └─metadata.json\n",
      "  │ └─model.py\n",
      "  └─config.pbtxt\n"
     ]
    }
   ],
   "source": [
    "sd.seedir(\n",
    "    export_path,\n",
    "    style=\"lines\",\n",
    "    itemlimit=10,\n",
    "    depthlimit=5,\n",
    "    exclude_folders=['.ipynb_checkpoints', '__pycache__'],\n",
    "    sort=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe7962cc-f26d-4a4a-b5a3-d214e0f37456",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Starting Triton Server"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c07c620-7d6c-4275-87fe-e5b94335bdb9",
   "metadata": {},
   "source": [
    "It is time to deploy all the models as an ensemble model to Triton Inference Serve [TIS](https://github.com/triton-inference-server). After we export the ensemble, we are ready to start the TIS. You can start triton server by using the following command on your terminal:\n",
    "\n",
    "```\n",
    "tritonserver --model-repository=/ensemble_export_path/ --backend-config=tensorflow,version=2\n",
    "```\n",
    "\n",
    "For the `--model-repository` argument, specify the same path as the `export_path` that you specified previously in the `ensemble.export` method. This command will launch the server and load all the models to the server. Once all the models are loaded successfully, you should see `READY` status printed out in the terminal for each loaded model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c0a798f-6abf-4cbb-87f8-f60a6e757092",
   "metadata": {},
   "source": [
    "### Retrieving Recommendations from Triton"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b0794b1-b9e0-4508-bf6e-cc823ac5c693",
   "metadata": {},
   "source": [
    "Once our models are successfully loaded to the TIS, we can now easily send a request to TIS and get a response for our query with `send_triton_request` utility function. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af9efbde-4dac-42f1-9ace-096f75bac2b5",
   "metadata": {},
   "source": [
    "Let's send a request to TIS for a given `user_id_raw` value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "d08a8975-9c32-467b-99ec-df66319f854b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   user_id_raw\n",
      "0            2\n"
     ]
    }
   ],
   "source": [
    "# read in data for request\n",
    "from merlin.core.dispatch import make_df\n",
    "\n",
    "# create a request to be sent to TIS\n",
    "request = make_df({\"user_id_raw\": [2]})\n",
    "request[\"user_id_raw\"] = request[\"user_id_raw\"].astype(np.int32)\n",
    "print(request)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28e9e27f-6658-4302-b142-08b05215e48f",
   "metadata": {},
   "source": [
    "Let's return raw item ids from TIS as top-k recommended items per given request."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "74ec62f2-5935-45c6-8058-e1cdade6f80f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ordered_ids': array([[214532372, 214531867, 214828618, 214530521, 214680037, 214823075,\n",
       "         214778081, 214612513, 214537616, 214745795]], dtype=int32),\n",
       " 'ordered_scores': array([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]], dtype=float32)}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = send_triton_request(request_schema, request, outputs)\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4605dbe-5f97-4b31-8ee4-ce7c1cb69d97",
   "metadata": {},
   "source": [
    "That's it! You finished deploying a multi-stage Recommender Systems on Triton Inference Server using the Merlin framework and the Milvus vector search library."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "merlin": {
   "containers": [
    "nvcr.io/nvidia/merlin/merlin-tensorflow-inference:latest"
   ]
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
